'''
https://leetcode.cn/problems/ones-and-zeroes/
'''
from collections import Counter
from functools import cache
from typing import List

import numpy as np


class Solution:
    # 意思就是，找出一个strs的子集：子集里的元素是s's
    # 这些s's number of 0's <= m and number of 1's <= n
    # return the max size of s's
    def findMaxForm(self, strs: List[str], m: int, n: int) -> int:
        counters = [Counter(s) for s in strs]

        # u:还能选多少0， v:还能选多少1, i:第i个单词
        @cache
        def f(i, u, v):
            if i == len(strs): return 0
            res = f(i + 1, u, v)
            if counters[i]['0'] <= u and counters[i]['1'] <= v:
                res = max(res, 1 + f(i + 1, u - counters[i]['0'], v - counters[i]['1']))
            return res

        return f(0, m, n)

    # 打表
    def findMaxForm2(self, strs: List[str], m: int, n: int) -> int:
        counters = [Counter(s) for s in strs]
        dp = np.zeros((len(strs) + 1, m + 1, n + 1), dtype=np.int64)
        # dp[len(strs),:,:] = 0
        for i in range(len(strs) - 1, -1, -1):
            for u in range(m + 1):
                for v in range(n + 1):
                    res = dp[i + 1, u, v]
                    if counters[i]['0'] <= u and counters[i]['1'] <= v:
                        res = max(res, 1 + dp[i + 1, u - counters[i]['0'], v - counters[i]['1']])
                    dp[i][u][v] = res
        # print(dp)
        return int(dp[0, m, n])

    # 打表原生数组
    def findMaxForm3(self, strs: List[str], m: int, n: int) -> int:
        counters = [Counter(s) for s in strs]
        dp = [[[0] * (n + 1) for _ in range(m + 1)] for _ in range(len(strs) + 1)]
        for i in range(len(strs) - 1, -1, -1):
            for u in range(m + 1):
                for v in range(n + 1):
                    res = dp[i + 1][u][v]
                    if counters[i]['0'] <= u and counters[i]['1'] <= v:
                        res = max(res, 1 + dp[i + 1][u - counters[i]['0']][v - counters[i]['1']])
                    dp[i][u][v] = res
        # print(dp)
        return dp[0][m][n]
    # 状态压缩
    def findMaxForm4(self, strs: List[str], m: int, n: int) -> int:
        counters = [Counter(s) for s in strs]
        A = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(len(strs) - 1, -1, -1):
            B = [[0] * (n + 1) for _ in range(m + 1)]
            for u in range(m + 1):
                for v in range(n + 1):
                    res = A[u][v]
                    if counters[i]['0'] <= u and counters[i]['1'] <= v:
                        res = max(res, 1 + A[u - counters[i]['0']][v - counters[i]['1']])
                    B[u][v] = res
            A = B
        return A[m][n]



strs = ["101000000", "1100001010", "11101000", "011010110", "0010001", "0011", "0111101111"]
m = 10
n = 11
print(Solution().findMaxForm3(strs, m, n))
print(Solution().findMaxForm4(strs, m, n))
